# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import enum
import json
import logging
import functools
import time
import hmac
import hashlib
import sys

from tornado import gen
from tornado.concurrent import Future
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect, WebSocketClosedError

from coin.base.timestamp import get_timestamp
from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin.exchange.base.rest.task_scheduler import Scheduler

BITFLYER_WS_API_HOST = 'ws.lightstream.bitflyer.com'


class OperatorType(enum.Enum):
  AUTH = 'auth'
  SUB = 'subscribe'
  UNSUB = 'unsubscribe'


# currently produces order event only
class BitflyerWsPrivateClient(object):
  def __init__(
      self,
      key_file,
      *,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=20,
      api_host=None,
      disable_ssl=False,
  ):
    self._ioloop = ioloop or IOLoop.current()
    scheme = 'wss://'
    api_host = api_host or BITFLYER_WS_API_HOST
    self._api_url = scheme + api_host + '/json-rpc'
    self._ws = None
    self._ready = False
    self._timeout = timeout_sec * 1e9
    self._num_timeouts = 0
    self._key = AuthKey.from_file(key_file)

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_futures = {}
    self._subscribed_topics = set()
    self._request_id = int(time.time() * 1000)

    self._stopped = False
    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)
    self._logger = logging.getLogger(__name__)

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _get_next_request_id(self):
    self._request_id += 1
    return str(self._request_id)

  @gen.coroutine
  def _check_future_timeout(self):
    timestamp = get_timestamp()
    for req_id, fut in self._pending_futures.copy().items():
      if timestamp - fut.timestamp > self._timeout:
        self._num_timeouts += 1
        fut = self._pending_futures.pop(req_id)
        fut.set_exception(TimeoutError('Websocket request timeout! (%s, %s)' %
                                       (req_id, fut.method)))

  @gen.coroutine
  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return
    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')

  def _gen_auth_msg(self):
    timestamp = int(round(time.time() * 1000))
    nonce = ''.join(reversed(str(timestamp * 1000)))
    message = f"{timestamp}{nonce}"
    signature = hmac.new(bytes(self._key.api_secret, 'utf8'),
                         bytes(message, 'utf8'),
                         digestmod=hashlib.sha256).hexdigest()
    params = dict()
    params['params'] = {
        'api_key': self._key.access_key,
        'timestamp': timestamp,
        'nonce': nonce,
        'signature': signature
    }
    return params

  def send_request(self, operation, *, topic=None, is_notification=False):
    if self._ws is None:
      self._logger.error('websocket not connected yet! operation=%s, params=%s',
                         operation.value,
                         topic)
      return

    fut = None
    request_id = None
    if not is_notification:
      fut = Future()
      fut.timestamp = get_timestamp()
      fut.operation = operation
      request_id = self._get_next_request_id()
      self._pending_futures[request_id] = fut

    request_msg = {'jsonrpc': '2.0', 'id': request_id}
    if operation == OperatorType.AUTH:
      request_msg['method'] = operation.value
      request_msg.update(self._gen_auth_msg())
    else:
      assert topic is not None
      request_msg['method'] = operation.value
      request_msg['params'] = {'channel': topic}

    self._logger.info('send request: %s\n', request_msg)
    self._ioloop.add_callback(self._write_message, request_msg)
    return fut

  def subscribe_topic(self, topic):
    operation = OperatorType.SUB
    fut = self.send_request(operation, topic=topic)
    return fut

  def unsubscribe_topic(self, topic):
    operation = OperatorType.UNSUB
    fut = self.send_request(operation, topic=topic)
    return fut

  @gen.coroutine
  def prepare_subscription_channels(self, channels):
    for channel in channels:
      response = yield self.subscribe_topic(channel)
      if 'error' in response:
        self.stop()
        return
      else:
        self._logger.info('Subscribe to bitflyer order status ch:%s', channel)

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket (%s).',
                        type(self).__name__,
                        self._api_url)
      self._ws = yield websocket_connect(self._api_url)
      self._ioloop.add_callback(self._run)

      self._logger.info('Send auth info!')
      response = yield self.send_request(operation=OperatorType.AUTH)
      if 'error' in response:
        self._logger.error('Authentication failed! %s', response)
        self.stop()
        return
      else:
        self._logger.info('Bitflyer WS API Successfully authenticated!')

      self._ready = True
      self._on_ready()
      self._logger.info('Bitflyer private client websocket connected.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    self._logger.info('Receive ws msg: %s\n', data)

    method = data.get('method', None)
    if method == 'channelMessage':
      self._notification_callback(data)
    elif method is None:  # something's repsonse. not subscription.
      req_id = data.get('id', None)
      if req_id is None:
        self._logger.error('No request id. %s', data)
        return

      operation = self._pending_futures[req_id].operation
      if operation == OperatorType.SUB:
        print("skip")
        # self._subscribed_topics.add(topic)
      fut = self._pending_futures.pop(req_id, None)
      if fut is None:
        self._logger.error('Unexpected message! %s', data)
      else:
        fut.set_result(data)

  @gen.coroutine
  def _ping(self):
    if not self._stopped:
      try:
        self._logger.info("ping (actually it's auth since no ping)")
        response = yield self.send_request(operation=OperatorType.AUTH)
      except Exception:
        self._logger.exception('Ping error')

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(15, self._ping)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Bitflyer ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    for topic in self._subscribed_topics:
      response = yield self.unsubscribe_topic(topic)
      self._logger.info('Topic (%s) unsubscribed!\n%s', topic, response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    yield gen.sleep(10)
    self._close_callback()


class TestBitflyerWsPrivateClient(object):
  def __init__(self, key_filepath, ioloop=None):
    self._key_filepath = key_filepath
    self._ioloop = ioloop
    self._ws_private_client = None

  def _prepare_ws_private_client(self, start=False):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    self._ws_private_client = BitflyerWsPrivateClient(
        self._key_filepath,
        ioloop=self._ioloop,
        close_callback=functools.partial(self._prepare_ws_private_client, True),
        on_ready=self._on_ws_private_client_ready,
    )

    if start:
      self._ws_private_client.start()

  def _on_ws_private_client_ready(self):
    channels = ['child_order_events']
    self._ws_private_client.prepare_subscription_channels(channels)

  def start(self):
    self._prepare_ws_private_client(start=True)


def main(argv):
  key_filepath = argv[1]  # '/home/taekwon/bitflyer_api_key.json'

  ioloop = IOLoop.current()
  test_client = TestBitflyerWsPrivateClient(key_filepath, ioloop)
  test_client.start()
  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
